package com.ikas.ai.framework.utils;

import com.google.common.collect.Maps;
import com.ikas.ai.framework.core.model.BatchDataFunction;
import com.ikas.ai.framework.core.model.PageBatchOutput;
import com.ikas.ai.framework.core.model.PageQueryInput;
import com.ikas.ai.framework.statistics.StatisticType;
import com.ikas.ai.framework.statistics.function.*;
import com.ikas.ai.utils.CollectionUtil;

import java.util.List;
import java.util.Map;

/**
 * @Description: 使用分页加载数据
 * @Author: qi.yilai
 * @Date: 2023/10/30 14:21
 * @Version: 1.0.0
 **/
public class LoadDataUtils {
    private static Map<StatisticType, CalcFunction> calcFunctionMap = Maps.newHashMap();

    private static final int DEFAULT_PAGE_SIZE = 30000;

    static {
        calcFunctionMap.put(StatisticType.COUNT, new CountCalcFunction());
        calcFunctionMap.put(StatisticType.NULL_COUNT, new NullCountCalcFunction());
        calcFunctionMap.put(StatisticType.NOT_NULL_COUNT, new NotNullCountCalcFunction());
        calcFunctionMap.put(StatisticType.NULL_RATIO, new NullRatioCalcFunction());
        calcFunctionMap.put(StatisticType.NOT_NULL_RATIO, new NotNullRatioCalcFunction());
        calcFunctionMap.put(StatisticType.SUM, new SumCalcFunction());

        calcFunctionMap.put(StatisticType.MAX, new MaxCalcFunction());
        calcFunctionMap.put(StatisticType.MIN, new MinCalcFunction());
        calcFunctionMap.put(StatisticType.AVG, new AvgCalcFunction());
        calcFunctionMap.put(StatisticType.STANDARD_DEVIATION, new StandardDeviationCalcFunction());
        calcFunctionMap.put(StatisticType.MEDIAN, new MedianCalcFunction());
    }

    /**
     * 获取所有的数据
     *
     * @param batchDataFunction
     * @return
     */
    public static Map<String, List<String>> loadAllData(BatchDataFunction batchDataFunction) {
        Map<String, List<String>> dataMap = Maps.newHashMap();
        int queryType = batchDataFunction.getQueryType();
        PageQueryInput pageQueryInput = createQuery(batchDataFunction);
        // 1=按时长取数据, 3=按时间段取数
        if (queryType == 1 || queryType == 3) {
            int page = 1;
            while(true) {
                pageQueryInput.setPage(page);
                pageQueryInput.setPageSize(DEFAULT_PAGE_SIZE);
                PageBatchOutput batchOutput = batchDataFunction.getFunction().apply(pageQueryInput);
                merge(batchOutput.getDataMap(), dataMap);
                if (!batchOutput.toBeContinue()) {
                    break;
                }
                page ++;
            }
        } else {
            // 2=按数据量取数
            PageBatchOutput batchOutput = batchDataFunction.getFunction().apply(pageQueryInput);
            dataMap = batchOutput.getDataMap();
        }
        return dataMap;
    }

    private static void merge(Map<String, List<String>> fromDataMap, Map<String, List<String>> toDataMap) {
        fromDataMap.entrySet().stream().forEach(entry -> {
            List<String> list = toDataMap.get(entry.getKey());
            if (CollectionUtil.isEmpty(list)) {
                toDataMap.put(entry.getKey(), entry.getValue());
            } else {
                list.addAll(entry.getValue());
            }
        });
    }

    /**
     * 分页拉取数据
     *
     * @param batchDataFunction
     * @param page
     * @return
     */
    public static PageBatchOutput loadPageData(BatchDataFunction batchDataFunction, int page) {
        int queryType = batchDataFunction.getQueryType();
        PageQueryInput pageQueryInput = createQuery(batchDataFunction);
        // 1=按时长取数据, 3=按时间段取数
        if (queryType == 1 || queryType == 2) {
            pageQueryInput.setPage(page);
            pageQueryInput.setPageSize(DEFAULT_PAGE_SIZE);
            PageBatchOutput batchOutput = batchDataFunction.getFunction().apply(pageQueryInput);
            return batchOutput;
        } else {
            throw new UnsupportedOperationException("【按数据量取数】不支持分页取数");
        }
    }

    private static PageQueryInput createQuery(BatchDataFunction batchDataFunction) {
        PageQueryInput pageQueryInput = new PageQueryInput();
        pageQueryInput.setQueryConfig(batchDataFunction.getQueryConfig());
        pageQueryInput.setMeteData(batchDataFunction.getMeteData());
        pageQueryInput.setQueryType(batchDataFunction.getQueryType());
        return pageQueryInput;
    }

    public static <T extends CalcFunction> T getFunc(StatisticType statisticType, Class<T> clazz) {
        CalcFunction function = calcFunctionMap.get(statisticType);
        if (function.getClass().isAssignableFrom(clazz)) {
            return (T)function;
        } else {
            throw new UnsupportedOperationException(String.format("暂不支持改%s统计方法 or clazz参数错误", statisticType.getMethod()));
        }
    }
}
